import numpy as np
import matplotlib.pyplot as plt
import math
import sympy

# 绘制等高线图
delta = 0.025
x = np.arange(-5.0, 5.0, delta)
y = np.arange(-5.0, 5.0, delta)
X, Y = np.meshgrid(x, y)
Z = (3 * X ** 2 + 2 * Y ** 2 + X * Y)

fig, ax = plt.subplots(figsize=(10, 10))
CS = ax.contour(X, Y, Z)
ax.clabel(CS, inline=True, fontsize=10)

# 目标函数
x_sym, y_sym, alpha = sympy.symbols("x y alpha", real=True)
origin_func = 3 * x_sym ** 2 + 2 * y_sym ** 2 + x_sym * y_sym
grad_x = sympy.diff(origin_func, x_sym)
grad_y = sympy.diff(origin_func, y_sym)
f_alpha = 3 * (x_sym + alpha * grad_x) ** 2 + 2 * (y_sym + alpha * grad_y) ** 2 + (x_sym + alpha * grad_x) * (
        y_sym + alpha * grad_y)
f_alpha_diff = sympy.diff(f_alpha, alpha)

ax.set_title('Gradient Descent\n' + str(origin_func), fontsize=15)

print("f(x, y) = " + str(origin_func))
print("f(x, y)对x求偏导：" + str(grad_x))
print("f(x, y)对y求偏导：" + str(grad_y))
print("f(alpha) =", f_alpha)
print("f(f_alpha_diff) =", f_alpha_diff, end="\n\n")

x0, y0 = 5.0, 5.0
print("初始点位置(%f, %f)" % (x0, y0))

# 误差
EPSILON = 0.01

x_set = [x0]
y_set = [y0]
i = 1
while True:
    dx = float(grad_x.evalf(subs={x_sym: x0, y_sym: y0}))
    dy = float(grad_y.evalf(subs={x_sym: x0, y_sym: y0}))

    # 根据梯度判断是否满足误差
    if math.sqrt(dx ** 2 + dy ** 2) <= EPSILON:
        break

    # 计算最优步长
    f = f_alpha_diff.evalf(subs={x_sym: x0, y_sym: y0})
    k = sympy.solve(f, alpha)[0]
    x0 += k * dx
    y0 += k * dy
    print("迭代%d轮" % (i), x0, y0)
    i += 1
    x_set.append(x0)
    y_set.append(y0)

ax.plot(x_set, y_set, c='r', marker='o')
plt.show()
